package mr

import (
	"fmt"
	"log"
	"net"
	"net/http"
	"net/rpc"
	"os"
	"sync"
	"time"
)

const WAITING_STATUS = 0
const DOING_STATUS = 1
const SUCCESS_STATUS = 2

type TaskStatus struct {
	lastChangeTime time.Time
	status         int
}

type Coordinator struct {
	// Your definitions here.
	mapTaskStatus    map[string]*TaskStatus
	reduceTaskStatus map[string]*TaskStatus
	nReduce          int
	nowTime          time.Time
	cLock            sync.Mutex
}

func (c *Coordinator) UpdateLastTaskStatus(lastTask *RpcTask) {
	if lastTask.TaskType == MAP_TASK {
		//fmt.Printf("finish map task:%v\n", lastTask.SourceFile)
		nowStatus := c.mapTaskStatus[lastTask.SourceFile]
		nowStatus.status = SUCCESS_STATUS
	} else if lastTask.TaskType == REDUCE_TASK {
		//fmt.Printf("finish reduce task:%v\n", lastTask.SourceFile)
		nowStatus := c.reduceTaskStatus[lastTask.SourceFile]
		nowStatus.status = SUCCESS_STATUS
	} else if lastTask.TaskType == INIT_TASK {
		//do nothing
	} else if lastTask.TaskType == WAIT_NEXT_TASK {
		//do nothing
	} else {
		fmt.Printf("error task complete\n")
	}
}

func (c *Coordinator) TryGetNextTask(
	taskMap *map[string]*TaskStatus, taskType int, newTask *RpcTask) (bool, bool) {
	allFinish := true
	for fileName, taskStatus := range *taskMap {
		if taskStatus.status != SUCCESS_STATUS {
			//fmt.Printf("%v not finish\n", fileName)
			allFinish = false
		}
		if taskStatus.status == WAITING_STATUS {
			newTask.TaskType = taskType
			newTask.SourceFile = fileName
			newTask.NReduce = c.nReduce
			taskStatus.status = DOING_STATUS
			taskStatus.lastChangeTime = c.nowTime
			//fmt.Printf("send task %v\n", fileName)
			return true, allFinish
		} else if taskStatus.status == DOING_STATUS &&
			c.nowTime.Sub(taskStatus.lastChangeTime).Seconds() > 10 {
			//fmt.Printf("reset task %v\n", fileName)
			//fmt.Printf("diff: %v\n", c.nowTime.Unix()-taskStatus.lastChangeTime.Unix())
			taskStatus.lastChangeTime = c.nowTime
		}
	}
	return false, allFinish
}

func (c *Coordinator) TryGetMapTask(newTask *RpcTask) (bool, bool) {
	return c.TryGetNextTask(&c.mapTaskStatus, MAP_TASK, newTask)
}

func (c *Coordinator) TryGetReduceTask(newTask *RpcTask) (bool, bool) {
	return c.TryGetNextTask(&c.reduceTaskStatus, REDUCE_TASK, newTask)
}

// Your code here -- RPC handlers for the worker to call.
func (c *Coordinator) GetNewTask(lastTask *RpcTask, newTask *RpcTask) error {
	fmt.Printf("recv lastTask: %v, newTask: %v\n", lastTask, newTask)
	defer fmt.Printf("send lastTask: %v, newTask: %v\n", lastTask, newTask)
	//fmt.Println("get new task start")
	c.cLock.Lock()
	defer c.cLock.Unlock()
	c.nowTime = time.Now()
	if lastTask != nil {
		c.UpdateLastTaskStatus(lastTask)
	}

	ok, allFinish := c.TryGetMapTask(newTask)
	if ok {
		//fmt.Printf("new map task: %v\n", newTask.SourceFile)
		return nil
	}
	if !allFinish {
		newTask.TaskType = WAIT_NEXT_TASK
		//fmt.Printf("new waiting task for map\n")
		return nil
	}
	fmt.Println("all map task finish")

	ok, allFinish = c.TryGetReduceTask(newTask)
	if ok {
		//fmt.Printf("new reduce task: %v\n", newTask.SourceFile)
		return nil
	}
	if !allFinish {
		newTask.TaskType = WAIT_NEXT_TASK
		//fmt.Printf("new waiting task for reduce\n")
		return nil
	}
	//fmt.Println("all reduce task finish")

	newTask.TaskType = STOP_WORK_TASK
	//fmt.Printf("new stop task\n")
	return nil
}

// start a thread that listens for RPCs from worker.go
func (c *Coordinator) server() {
	rpc.Register(c)
	rpc.HandleHTTP()
	//l, e := net.Listen("tcp", ":1234")
	sockname := coordinatorSock()
	os.Remove(sockname)
	l, e := net.Listen("unix", sockname)
	if e != nil {
		log.Fatal("listen error:", e)
	}
	go http.Serve(l, nil)
}

// main/mrcoordinator.go calls Done() periodically to find out
// if the entire job has finished.
func (c *Coordinator) Done() bool {
	// Your code here.
	c.cLock.Lock()
	defer c.cLock.Unlock()
	for _, value := range c.mapTaskStatus {
		if value.status != SUCCESS_STATUS {
			return false
		}
	}
	for _, value := range c.reduceTaskStatus {
		if value.status != SUCCESS_STATUS {
			return false
		}
	}
	return true
}

// create a Coordinator.
// main/mrcoordinator.go calls this function.
// nReduce is the number of reduce tasks to use.
func MakeCoordinator(files []string, nReduce int) *Coordinator {
	c := Coordinator{}

	// Your code here.
	c.mapTaskStatus = make(map[string]*TaskStatus)
	c.reduceTaskStatus = make(map[string]*TaskStatus)
	for _, filePath := range files {
		nowStatus := TaskStatus{
			status:         WAITING_STATUS,
			lastChangeTime: time.Now(),
		}
		c.mapTaskStatus[filePath] = &nowStatus
	}

	for i := 0; i < nReduce; i++ {
		filePath := fmt.Sprintf("%v", i)
		nowStatus := TaskStatus{
			status:         WAITING_STATUS,
			lastChangeTime: time.Now(),
		}
		c.reduceTaskStatus[filePath] = &nowStatus
	}

	//for fileName := range c.mapTaskStatus {
	//	fmt.Printf("add map task: %v\n", fileName)
	//}
	//for fileName := range c.reduceTaskStatus {
	//	fmt.Printf("add reduce task: %v\n", fileName)
	//}

	c.nReduce = nReduce

	c.server()
	return &c
}
